import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os
import json
from tqdm import tqdm
# GPU setting
GPU_ID = '1'
os.environ['CUDA_VISIBLE_DEVICES'] = GPU_ID
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# ckpt setting
CKPT_PATH = 'ckpt/swin_model_100_0.977.pth'

# submit_file
SUBMIT = './submit_json'
submit_file = SUBMIT + '/swin_transformer.json'
if not os.path.exists(SUBMIT):
    os.mkdir(SUBMIT)

# test data path
PATH = './data4test'

# preprocessing
classes = ('GuideSign', 'M1', 'M4', 'M5', 'M6', 'M7', 'P1', 'P10_50', 'P12', 'W1')
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# create and load model
model = torch.load(CKPT_PATH)
model.eval()
model.to(DEVICE)

# testing
# saving as json file
testList = os.listdir(PATH)
print(testList)
print(type(testList))
pred_list = []
with open(submit_file, 'w') as f:
    for file in tqdm(testList):
        img=Image.open(PATH + '/' + file)
        img=transform_test(img)
        img.unsqueeze_(0)
        img = Variable(img).to(DEVICE)
        out=model(img)
        # Predict
        _, pred = torch.max(out.data, 1)
        # print('Image Name:{}, predict:{}'.format('test_dataset/'+file, int(pred)))
        pred_list.append({"filename" : 'test_dataset/'+file, "label" : int(pred)})
    print('predction Done! \n Generating json file!')
    pred_dict = {"annotations" : pred_list}
    f.write(json.dumps(pred_dict))
print(f'json file had generated saving as {submit_file}')